from __future__ import annotations

from typing import Optional, Tuple
from ppsci.equation import PDE
import sympy as sp


class Burgers(PDE):
    r"""Class for Burgers' equation in 1D, 2D, or 3D.

    $
    \mathb{u}_t + \mathb{u} \cdot \nabla \mathb{u} - \nu \nabla^2 \mathb{u} = 0
    $

    Args:
        dim (int): Spatial dimension (1, 2, or 3).
        time (bool): Whether to include time dimension.
        nu (float): Viscosity.
        detach_keys (Optional[Tuple[str, ...]]): Keys used for detach during computing.
            Defaults to None.

    Examples:
        >>> import ppsci
        >>> pde = ppsci.equation.Burgers(2, True, 1.0)
    """

    def __init__(
        self,
        dim: int,
        time: bool = True,
        nu: float = 1.0,
        detach_keys: Optional[Tuple[str, ...]] = None,
    ):
        super().__init__()
        self.detach_keys = detach_keys
        self.time = time
        self.dim = dim

        # Create symbols for time and space
        t, x, y, z = self.create_symbols("t x y z")
        if dim == 1:
            invars = (x,)
        elif dim == 2:
            invars = (x, y)
        elif dim == 3:
            invars = (x, y, z)
        else:
            raise ValueError("Dimension should be 1, 2, or 3.")

        if time:
            invars = (t,) + invars

        ##from NS equation
        # Create functions for each dimension
        u = self.create_function("u1", invars)
        v = self.create_function("u2", invars) if dim >= 2 else sp.Number(0)
        w = self.create_function("u3", invars) if dim >= 3 else sp.Number(0)

        momentum_x = (
            u.diff(t)
            + u * u.diff(x)
            + v * u.diff(y)
            + w * u.diff(z)
            - (
                (nu * u.diff(x)).diff(x)
                + (nu * u.diff(y)).diff(y)
                + (nu * u.diff(z)).diff(z)
            )
        )
        momentum_y = (
            v.diff(t)
            + u * v.diff(x)
            + v * v.diff(y)
            + w * v.diff(z)
            - (
                (nu * v.diff(x)).diff(x)
                + (nu * v.diff(y)).diff(y)
                + (nu * v.diff(z)).diff(z)
            )
        )
        momentum_z = (
            w.diff(t)
            + u * w.diff(x)
            + v * w.diff(y)
            + w * w.diff(z)
            - (
                (nu * w.diff(x)).diff(x)
                + (nu * w.diff(y)).diff(y)
                + (nu * w.diff(z)).diff(z)
            )
        )
        self.add_equation("burgers_eq1", momentum_x)
        if self.dim == 2:
            self.add_equation("burgers_eq2", momentum_y)
        elif self.dim == 3:
            self.add_equation("burgers_eq2", momentum_y)
            self.add_equation("burgers_eq3", momentum_z)
